import torch
from torch.nn import functional as F
import numpy as np
import matplotlib.pyplot as plt


def calculate_distances(tensor):
    tensor = tensor.squeeze(0)
    h, w = tensor.shape[1:]

    up_distances = torch.zeros((h, w), dtype=torch.float32)
    up_similarities = F.cosine_similarity(tensor[:, 1:, :], tensor[:, :-1, :], dim=0)
    up_distances[1:, :] = 1 - up_similarities  # 余弦距离 = 1 - 余弦相似度

    left_distances = torch.zeros((h, w), dtype=torch.float32)
    left_similarities = F.cosine_similarity(tensor[:, :, 1:], tensor[:, :, :-1], dim=0)
    left_distances[:, 1:] = 1 - left_similarities

    top_left_distances = torch.zeros((h, w), dtype=torch.float32)
    top_left_similarities = F.cosine_similarity(tensor[:, 1:, 1:], tensor[:, :-1, :-1], dim=0)
    top_left_distances[1:, 1:] = 1 - top_left_similarities

    return up_distances, left_distances, top_left_distances


def threshold_distances(up_distances, left_distances, top_left_distances, threshold):
    up_mask = up_distances < threshold
    left_mask = left_distances < threshold
    top_left_mask = top_left_distances < threshold
    return up_mask, left_mask, top_left_mask


def flood_fill(mask):
    h, w = mask.shape
    labels = torch.zeros((h, w), dtype=torch.int32)
    label = 1

    for i in range(h):
        for j in range(w):
            if mask[i, j] and labels[i, j] == 0:
                stack = [(i, j)]
                while stack:
                    x, y = stack.pop()
                    if 0 <= x < h and 0 <= y < w and mask[x, y] and labels[x, y] == 0:
                        labels[x, y] = label
                        stack.extend([(x - 1, y), (x + 1, y), (x, y - 1), (x, y + 1)])
                label += 1

    return labels


def visualize_and_save(labels, output_path):
    num_labels = int(labels.max()) + 1
    colors = np.random.rand(num_labels, 3)
    colors[0] = [0, 0, 0]

    colored_image = colors[labels.numpy()]

    plt.imsave(output_path, colored_image)


def find_threshold_for_categories(tensor, target_categories, max_iterations=100, tolerance=1):
    # 初始化二分查找的范围
    low_threshold = 0.0
    high_threshold = 1.0

    for _ in range(max_iterations):
        # 计算中间的 threshold
        mid_threshold = (low_threshold + high_threshold) / 2

        # 计算距离和掩码
        up_distances, left_distances, top_left_distances = calculate_distances(tensor)
        up_mask, left_mask, top_left_mask = threshold_distances(up_distances, left_distances, top_left_distances,
                                                                mid_threshold)
        combined_mask = up_mask | left_mask | top_left_mask

        # 进行洪水填充
        combined_labels = flood_fill(combined_mask)
        num_categories = int(combined_labels.max()) + 1

        # 判断是否满足条件
        if abs(num_categories - target_categories) <= tolerance:
            return mid_threshold
        elif num_categories > target_categories:
            high_threshold = mid_threshold
        else:
            low_threshold = mid_threshold
        if high_threshold-low_threshold<1e-2:
            return None
    return None


# 加载数据
tensor = torch.load('image_tensor.pt')
tensor = tensor.to('cuda')
tensor=torch.ones([3, 100, 100])
# 指定目标类别数量
target_categories = 1

# 寻找合适的 threshold
threshold = find_threshold_for_categories(tensor, target_categories)

if threshold is not None:
    print(f"找到合适的 threshold: {threshold}")

    up_distances, left_distances, top_left_distances = calculate_distances(tensor)
    up_mask, left_mask, top_left_mask = threshold_distances(up_distances, left_distances, top_left_distances, threshold)
    combined_mask = up_mask | left_mask | top_left_mask

    combined_labels = flood_fill(combined_mask)

    # 统计每个类别的像素数量
    unique, counts = torch.unique(combined_labels, return_counts=True)
    for label, count in zip(unique, counts):
        print(f"类别 {label.item()} 所占的像素数量: {count.item()}")

    output_path = 'output_image.png'
    visualize_and_save(combined_labels, output_path)
    print(f"结果已保存为 {output_path}")
else:
    print("未找到合适的 threshold。")